package data.platform.store;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.datatype.jsr310.JavaTimeModule;
import data.platform.common.service.command.KeyValueRepository;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.rocksdb.RocksDB;
import org.rocksdb.RocksDBException;
import org.rocksdb.RocksIterator;
import org.rocksdb.WriteOptions;
import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
import org.springframework.stereotype.Repository;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;
import reactor.core.scheduler.Schedulers;
import reactor.util.function.Tuple2;
import reactor.util.function.Tuples;

import javax.annotation.PostConstruct;
import java.util.ArrayList;
import java.util.List;

@ConditionalOnBean(name = "rocksDBConfig")
@Slf4j
@Repository
@RequiredArgsConstructor
public class RocksDBRepositoryImpl implements KeyValueRepository<String, byte[]> {

    final WriteOptions writeOptions;

    final RocksDB rocksDB;

    private ObjectMapper objectMapper;

    @PostConstruct
    public void init() {
        objectMapper = new ObjectMapper();
        objectMapper.registerModule(new JavaTimeModule());
    }

    @Override
    public void put(String key, byte[] value) {
        try {
            rocksDB.put(writeOptions, key.getBytes(), value);
        } catch (Exception e) {
            log.error("Error saving entry in RocksDB, cause: {}, message: {}", e.getCause(), e.getMessage());
        }
    }

    @Override
    public Mono<byte[]> get(String key) {
        return Mono.fromCallable(() -> {
            try {
                return rocksDB.get(key.getBytes());
            } catch (Exception e) {
                log.error("Error retrieving the entry in RocksDB from key: {}, cause: {}, message: {}", key, e.getCause(), e.getMessage());
            }
            return null;
        }).subscribeOn(Schedulers.boundedElastic());
    }

    @Override
    public void delete(String key) {
        try {
            rocksDB.delete(writeOptions, key.getBytes());
        } catch (RocksDBException e) {
            log.error("Error deleting entry in RocksDB, cause: {}, message: {}", e.getCause(), e.getMessage());
        }
    }

    @Override
    public Flux<Tuple2<String, byte[]>> getByPrefixKey(String prefixKey) {
        List<Tuple2<String, byte[]>> values = new ArrayList<>();
        RocksIterator rocksIterator = rocksDB.newIterator();
        for (rocksIterator.seek(prefixKey.getBytes()); rocksIterator.isValid(); rocksIterator.next()) {
            String key = new String(rocksIterator.key());
            if (key.startsWith(prefixKey)) {
                try {
                    values.add(Tuples.of(key, rocksIterator.value()));
                } catch (Exception e) {
                    log.error("Error retrieving the entry in RocksDB from key: {}, cause: {}, message: {}", key, e.getCause(), e.getMessage());
                }
            }
        }
        return Flux.fromIterable(values);
    }

}
